"""Cloud detection methods based on simple per-pixel heuristics

This module hosts a set of simple, quick yet efficient per-pixel heuristics based cloud detection methods.
Each method should accept a given ndarray with a dimension of [n_bands X n_rows X n_cols] representing the image,
and a `prob` criteria value for generating the binary mask.
Additional keyword arguments could be used to further tune the method's behavior.
The result should be a ndarray, representing cloud probability or binary cloud mask.

Currently supported methods are:
1. `detect_cloud_hsv`

We are going to support more methods, including:

1. `detect_cloud_dcp`
2. `detect_cloud_red_vs_blue`

"""
from typing import Optional
import numpy as np
import numba as nb

from common.typing_utils import numeric


def detect_cloud_hsv(
    nda_img: np.ndarray,
    prob: Optional[numeric] = 0.5,
    internal_threshold: numeric = 0.1,
    internal_spread: numeric = 4.0,
) -> np.ndarray:
    """Detect cloud in a given remote sensing image based on the HSV threshold method.

    Parameters
    ----------
    nda_img : np.ndarray
        The ndarray with a dimension of [n_bands X n_rows X n_cols] for the image.
        It is expected that `n_bands` is greater than or equal to 3.
        The first 3 bands are considered to be R-G-B bands in further computation.
    prob : Optional[numeric], optional
        A criteria value on detected cloud probability,
        should be None, indicating returning the cloud probability,
        or a number within the range of 0 - 1.0, indicating returning a binary mask after applying the criteria.
        By default 0.5.
    internal_threshold: numeric, optional
        Internal threshold value to be used, by default 0.1. See Notes for more details.
    logistic_spread: numeric
        Internal spread value to be used, by default 4. See Notes for more details.

    Returns
    -------
    np.ndarray
        Result ndarray for the detected cloud probability or binary mask.

    Notes
    -----
    This method detects cloud based on the simple heuristic of brightness (or value v) minus saturation (s).
    If V-S is greater than a given threshold, it is considered as cloud.
    We further convert the V-S with a logistic function $P = \\frac{1}{1 + \\exp{spread * (-x + threshold)}}$.
    By default, the `internal_threshold` is 0.1 and the `internal_spread` is 4.
    The combination of those two would result in a V-S of 0.1 ends up with a cloud probability of 0.5,
    and a V-S of 0.8 ends up with a cloud probability of 0.94.
    For an end user, you may mainly change the `prob` to either get the cloud probability map,
    or tune the probability criteria.
    If you are familiar with the brightness-vs-saturation heuristics, you may further tune the two internal args.

    References
    ----------
    Qingsong Zhu, Jiaming Mai, and Ling Shao. A fast single image haze removal algorithm using color attenuation prior.
        IEEE Transactions on Image Processing, 24(11):3522–3533, 2015. 2
    """
    if (len(nda_img.shape) != 3) or (nda_img.shape[0] < 3):
        raise ValueError(
            "The given `nda_img` must be of shape [n_bands X n_rows X n_cols] with more than 3 bands."
        )

    nda_max = np.max(nda_img[:3, :, :], axis=0)
    nda_min = np.min(nda_img[:3, :, :], axis=0)
    nda_delta = nda_max - nda_min
    v_channel = nda_max
    s_channel = nda_delta / v_channel
    s_channel[v_channel == 0] = 0
    v_s = v_channel - s_channel

    # prob = 1 / (1 + exp(spread * (-x + threshold)))
    v_s *= -1
    v_s += internal_threshold
    v_s *= internal_spread
    v_s = np.exp(v_s, out=v_s)
    v_s += 1
    v_s = np.divide(1, v_s, out=v_s)

    if prob is not None:
        return v_s > prob
    else:
        return v_s


@nb.jit(nopython=True)
def _compute_rgb_dark_channel_index_nb(img: np.ndarray, patch_size: int = 15):
    """Compute the average of darkest value in all patches.

    Parameters
    -----------
    img: numpy.ndarray
        Input image as an UInt8 numpy array.
        - Bands: [R, G, B] or [B, G, R]
        - Pixel DN range: [0, 255].
        - Shape: (n_bands, n_rows, n_cols)
    patch_size: int
        Patch size for calculation, Default: 15

    Returns
    -------
    float
        average of darkest value
    """
    max_val = 255
    count = 0
    sum_ = 0
    num_row, num_col = img.shape[1], img.shape[2]
    for row in range(0, num_row, patch_size):
        for col in range(0, num_col, patch_size):
            img_patch = img[:, row : row + patch_size, col : col + patch_size]
            if not np.any(img_patch):
                continue

            min_in_patch = max_val
            for val in img_patch.ravel():
                if val > 0:
                    if val < min_in_patch:
                        min_in_patch = val
            # if min_in_patch < max_val:
            sum_ += min_in_patch
            count += 1
    return float(sum_) / (count + 1e-8)


def compute_dark_channel_index(img: np.ndarray, patch_size: int = 15, pxl_max=4000):
    """Compute the average of darkest value in all patches.

    Parameters
    -----------
    img: numpy.ndarray
        Input image as an UInt16 numpy array, similar to Sentinel-2 raw data.
        - Bands: [R, G, B] or [B, G, R]
        - Pixel DN range: (0, 10000). For a few pixels, it may exceed 10000.
        - Shape: (n_bands, n_rows, n_cols)
    patch_size: int
        Patch size for calculation, Default: 15
    pxl_max: int
        Max pixel value of the image. Default: 4000.

    Returns
    -------
    float
        average of darkest value
    """
    if (
        not isinstance(img, np.ndarray)
        or img.ndim != 3
        or img.dtype != np.uint16
        or img.shape[0] != 3
    ):
        raise ValueError(
            "Invalid input `img`. Only accepts an image of an UInt16 numpy array "
            "with shape `(n_bands, n_rows, n_cols)` and bands `(R,G,B)`."
        )

    if np.all(img == 0):
        raise ValueError("Cannot input all-zero `img` array.")

    # Normalize UInt16 remote sensing RGB imagery to UInt8 natural RGB imagery.
    max_val = 255
    img = np.clip(img, 0, pxl_max).astype(np.float32)
    img = (img / pxl_max * max_val).astype(np.uint8)
    dci = _compute_rgb_dark_channel_index_nb(img, patch_size=patch_size)
    return dci


def detect_cloud_red_vs_blue(
    nda_img: np.ndarray,
    prob: Optional[numeric] = 0.5,
) -> np.ndarray:
    """Cloud detection based on red-vs-blue characteristics"""
    pass
